from commands import add, admin
import commands
     
@admin
def vampmode(connection):
    if connection.protocol.vamp == True:
        connection.protocol.send_chat('Vampire Mode Disabled.')
        connection.protocol.vamp = False
    else:
        if connection.protocol.vamp is None:
            connection.protocol.vamp = set()
        connection.protocol.send_chat('Vampire Mode Enabled.')
        connection.protocol.vamp = True
add(vampmode)


def apply_script(protocol, connection, config):
    
    class VampConnection(connection):
        def on_hit(self, hit_amount, hit_player):
            if self.protocol.vamp == 1 and self.team is not hit_player.team:
                    self.hit_player.vamp_hit(hit_amount)
            return connection.on_hit(self, hit_amount, hit_player)

        def vamp_hit(self, hit_amount):
            self.set_hp(self.hp + hit_amount)

    return protocol, VampConnection